{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! /usr/bin/env python\n",
    "# coding=utf-8\n",
    "\n",
    "import cv2\n",
    "import time\n",
    "import numpy as np\n",
    "import core.utils as utils\n",
    "import tensorflow as tf\n",
    "from core.yolov3 import YOLOv3, decode\n",
    "\n",
    "\n",
    "video_path      = \"./docs/road.mp4\"\n",
    "num_classes     = 25\n",
    "input_size      = 416\n",
    "\n",
    "input_layer  = tf.keras.layers.Input([input_size, input_size, 3])\n",
    "feature_maps = YOLOv3(input_layer)\n",
    "\n",
    "bbox_tensors = []\n",
    "for i, fm in enumerate(feature_maps):\n",
    "    bbox_tensor = decode(fm, i)\n",
    "    bbox_tensors.append(bbox_tensor)\n",
    "\n",
    "model = tf.keras.Model(input_layer, bbox_tensors)\n",
    "# utils.load_weights(model, \"./yolov3.weights\")\n",
    "model.load_weights(\"./yolov3\")\n",
    "vid = cv2.VideoCapture(video_path)\n",
    "\n",
    "# vid = cv2.VideoCapture(0)\n",
    "while True:\n",
    "    return_value, frame = vid.read()\n",
    "    if return_value:\n",
    "        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "    else:\n",
    "        raise ValueError(\"No image!\")\n",
    "    frame_size = frame.shape[:2]\n",
    "    image_data = utils.image_preporcess(np.copy(frame), [input_size, input_size])\n",
    "    image_data = image_data[np.newaxis, ...].astype(np.float32)\n",
    "\n",
    "    prev_time = time.time()\n",
    "    pred_bbox = model.predict(image_data)\n",
    "    curr_time = time.time()\n",
    "    exec_time = curr_time - prev_time\n",
    "\n",
    "    pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]\n",
    "    pred_bbox = tf.concat(pred_bbox, axis=0)\n",
    "    bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size, 0.3)\n",
    "    bboxes = utils.nms(bboxes, 0.45, method='nms')\n",
    "    image = utils.draw_bbox(frame, bboxes)\n",
    "\n",
    "    result = np.asarray(image)\n",
    "    info = \"time: %.2f ms\" %(1000*exec_time)\n",
    "    cv2.putText(result, text=info, org=(50, 70), fontFace=cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                fontScale=1, color=(255, 0, 0), thickness=2)\n",
    "    cv2.namedWindow(\"result\", cv2.WINDOW_AUTOSIZE)\n",
    "    result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n",
    "    cv2.imshow(\"result\", result)\n",
    "    if cv2.waitKey(1) & 0xFF == ord('q'): break\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
